#! /home/handuser/miniconda3/envs/yjh/bin/python3
import os, sys
import numpy as np
import rospy
from sensor_msgs.msg import Image
import torchvision.transforms as transforms
import torch
import torch.nn.functional as F
import gorilla
import copy

HERE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(HERE_DIR, '../../utils/'))
from image_converter import convert_Image_to_nparray, convert_nparray_to_Image
sys.path.append(os.path.join(HERE_DIR, '../../instance_segmentation/'))
from scripts.detector import InstSeg

import cv2

DPDN_DIR = os.path.join(HERE_DIR, '../../objpose_estimation/dpdn')
sys.path.append(os.path.join(DPDN_DIR, 'provider'))
sys.path.append(os.path.join(DPDN_DIR, 'model'))
sys.path.append(os.path.join(DPDN_DIR, 'model', 'pointnet2'))
sys.path.append(os.path.join(DPDN_DIR, 'utils'))
from DPDN import Net as DPDN_Net
from data_utils import get_bbox

SPD_DIR = os.path.join(HERE_DIR, '../../objpose_estimation/spd')
sys.path.append(os.path.join(SPD_DIR, 'lib'))
from network import DeformNet
from align import estimateSimilarityTransform

sys.path.append(HERE_DIR)
from draw_utils import draw_pose_detections


class ObjPoseEst:
    def __init__(self, device) -> None:
        here = os.path.dirname(os.path.abspath(__file__))
        self.detector = InstSeg(device)
        self.device = device
        
        self.show_pc = False

        ckpt = os.path.join(here, '../assets/epoch_30.pth')
        self.n_pts = 1024
        self.mean_shapes = np.load(os.path.join(here, '../assets/mean_shapes.npy'))

        model = 'SPD'
        if model == 'DPDN':
            ckpt = os.path.join(here, '../assets/epoch_30.pth')
            self.objpose_model = DPDN_Net(nclass=6, nprior=1024)
            gorilla.solver.load_checkpoint(model=self.objpose_model, filename=ckpt)
            self.run = self.run_DPDN
        elif model == 'SPD':
            ckpt = os.path.join(here, '../assets/model_50.pth')
            self.objpose_model = DeformNet(n_cat=6, nv_prior=1024)
            self.objpose_model.load_state_dict(torch.load(ckpt))       
            self.run = self.run_SPD
        else:
            raise NotImplementedError
        self.objpose_model.eval()
        self.objpose_model.to(device)        

        self.cam_K = np.array([[603.1332397460938, 0, 323.1732177734375],
                               [0, 601.6819458007812, 241.5373077392578],
                               [0, 0, 1]])
        self.xmap = np.array([[i for i in range(640)] for j in range(480)])
        self.ymap = np.array([[j for i in range(640)] for j in range(480)])
        self.norm_scale = 1000.0
        self.img_size = 192
        self.norm_color = transforms.Compose(
                            [transforms.ToTensor(),
                            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])        

    def serve_as_regular_publisher(self, rate_hz=2):
        rospy.init_node('objpose_estimation', anonymous=True)
        self.puber1 = rospy.Publisher('/dexgrasp/mask', Image, queue_size=1)
        self.puber2 = rospy.Publisher('/dexgrasp/pose', Image, queue_size=1)
        rate = rospy.Rate(rate_hz)
        while not rospy.is_shutdown():
            img_raw = rospy.wait_for_message("/camera/color/image_raw", Image)
            depth_raw = rospy.wait_for_message("/camera/aligned_depth_to_color/image_raw", Image)
            result_instseg, img_masked_msg = self.detector.run(img_raw)
            self.puber1.publish(img_masked_msg)
            result_pose, img_posed_msg = self.run(img_raw, depth_raw, result_instseg)
            self.puber2.publish(img_posed_msg)
            rate.sleep()

    def prepare_cropped_data(self, img, depth, result_instseg):
        cam_fx, cam_fy = self.cam_K[0, 0], self.cam_K[1, 1]
        cam_cx, cam_cy = self.cam_K[0, 2], self.cam_K[1, 2]       

        if type(result_instseg['class_ids']) is int:
            for k, v in result_instseg.items():
                result_instseg[k] = np.expand_dims(np.array(v), axis=0)

        # prepare frame data
        f_points, f_rgb, f_choose, f_catId, f_prior = [], [], [], [], []
        valid_inst = []
        num_insts = len(result_instseg['class_ids'])

        for i in range(num_insts):
            cat_id = result_instseg['class_ids'][i]
            prior = self.mean_shapes[cat_id]
            rmin, rmax, cmin, cmax = get_bbox(result_instseg['bboxes'][i])
            rmin, rmax, cmin, cmax = int(rmin), int(rmax), int(cmin), int(cmax)
            mask = np.logical_and(result_instseg['masks'][i, :, :], depth > 0)
            choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
            
            if len(choose) < 32:
                continue
            else:
                valid_inst.append(i)

            # process objects with valid depth observation
            if len(choose) > self.n_pts:
                c_mask = np.zeros(len(choose), dtype=int)
                c_mask[:self.n_pts] = 1
                np.random.shuffle(c_mask)
                choose = choose[c_mask.nonzero()]
            else:
                choose = np.pad(choose, (0, self.n_pts-len(choose)), 'wrap')
            depth_masked = depth[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis]
            xmap_masked = self.xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis]
            ymap_masked = self.ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis]
            pt2 = depth_masked / self.norm_scale
            pt0 = (xmap_masked - cam_cx) * pt2 / cam_fx
            pt1 = (ymap_masked - cam_cy) * pt2 / cam_fy
            points = np.concatenate((pt0, pt1, pt2), axis=1)
            rgb = img[rmin:rmax, cmin:cmax, :]
            rgb = cv2.resize(rgb, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
            rgb = self.norm_color(rgb)
            crop_w = rmax - rmin
            ratio = self.img_size / crop_w
            col_idx = choose % crop_w
            row_idx = choose // crop_w
            choose = (np.floor(row_idx * ratio) * self.img_size + np.floor(col_idx * ratio)).astype(np.int64)
            # concatenate instances
            f_points.append(points)
            f_rgb.append(rgb)
            f_choose.append(choose)
            f_catId.append(cat_id)
            f_prior.append(prior)
        return valid_inst, [f_points, f_rgb, f_choose, f_catId, f_prior]

    def run_SPD(self, image_raw:Image, depth_raw:Image, result_instseg):
        if result_instseg is None:
            return None, image_raw
        
        img = convert_Image_to_nparray(image_raw)
        depth = convert_Image_to_nparray(depth_raw)

        valid_inst, [f_points, f_rgb, f_choose, f_catId, f_prior] = self.prepare_cropped_data(img, depth, result_instseg)

        num_valid_inst = len(valid_inst)
        if num_valid_inst == 0:
            return None, image_raw
        valid_inst_np = np.array(valid_inst)
        result_instseg['class_ids'] = result_instseg['class_ids'][valid_inst_np]
        result_instseg['bboxes'] = result_instseg['bboxes'][valid_inst_np]
        result_instseg['scores'] = result_instseg['scores'][valid_inst_np]
        f_sRT = np.zeros((num_valid_inst, 4, 4), dtype=float)
        f_s = np.zeros((num_valid_inst), dtype=float)
        f_R = np.zeros((num_valid_inst, 3, 3), dtype=float)
        f_t = np.zeros((num_valid_inst, 3), dtype=float)
        f_size = np.zeros((num_valid_inst, 3), dtype=float)

        if len(valid_inst):
            f_points = torch.cuda.FloatTensor(f_points)
            f_rgb = torch.stack(f_rgb, dim=0).cuda()
            f_choose = torch.cuda.LongTensor(f_choose)
            f_catId = torch.cuda.LongTensor(f_catId)
            f_prior = torch.cuda.FloatTensor(f_prior)
            # inference
            assign_mat, deltas = self.objpose_model(f_points, f_rgb, f_choose, f_catId, f_prior)
            inst_shape = f_prior + deltas
            assign_mat = F.softmax(assign_mat, dim=2)
            f_coords = torch.bmm(assign_mat, inst_shape)  # bs x n_pts x 3

            f_coords = f_coords.detach().cpu().numpy()
            f_points = f_points.cpu().numpy()
            f_choose = f_choose.cpu().numpy()
            f_insts = inst_shape.detach().cpu().numpy()

            # i: inverse transformation
            cuda_pred_i_sR = torch.zeros([num_valid_inst, 3, 3], dtype=torch.float32).cuda()
            cuda_pred_t = torch.zeros([num_valid_inst, 1, 3], dtype=torch.float32).cuda()

            for i in range(num_valid_inst):
                choose = f_choose[i]
                _, choose = np.unique(choose, return_index=True)
                nocs_coords = f_coords[i, choose, :]
                f_size[i] = 2 * np.amax(np.abs(f_insts[i]), axis=0)
                points = f_points[i, choose, :]
                _, _, _, pred_sRT = estimateSimilarityTransform(nocs_coords, points)
                if pred_sRT is None:
                    pred_sRT = np.identity(4, dtype=float)
                f_sRT[i] = pred_sRT

                s = np.cbrt(np.linalg.det(f_sRT[i][:3, :3]))
                R = f_sRT[i][:3, :3] / s
                i_sR = R / s

                f_s[i] = s
                f_R[i] = R
                f_t[i] = f_sRT[i][:3, 3]

                cuda_pred_i_sR[i] = torch.from_numpy(i_sR).float().cuda()
                cuda_pred_t[i] = torch.from_numpy(f_t[i]).view(1,3).float().cuda()

            cuda_obj_pcs = torch.bmm(torch.add(torch.cuda.FloatTensor(f_points), -cuda_pred_t), cuda_pred_i_sR)

            # filter point cloud in OBJCOORD
            obj_pcs = np.zeros_like(f_points)
            for i in range(cuda_obj_pcs.shape[0]):
                obj_pc = cuda_obj_pcs[i].detach().cpu().numpy()
                assert obj_pc.shape[1] == 3
                norm_pt = np.linalg.norm(obj_pc, axis=1)
                index = np.where(norm_pt < 0.5)
                obj_pc_filterd = obj_pc[index]
                if len(obj_pc_filterd) >= self.n_pts:
                    sample_index = np.random.choice(obj_pc_filterd.shape[0], self.n_pts, replace=False)
                else:
                    sample_index = np.random.choice(obj_pc_filterd.shape[0], self.n_pts, replace=True)
                    obj_pc_filterd = obj_pc_filterd[sample_index, :]
                obj_pcs[i] = obj_pc_filterd

                if self.show_pc:
                    import trimesh
                    pc = trimesh.PointCloud(obj_pc, colors=np.array([255,0,0,255]))
                    pc_filtered = trimesh.PointCloud(obj_pc_filterd)
                    pred_pc = trimesh.PointCloud(f_insts[i], colors=np.array([0,255,0,255]))
                    trimesh.Scene([pc, pc_filtered, pred_pc]).show()
        result = {}
        # 2D
        result['pred_class_ids'] = result_instseg['class_ids']
        result['pred_bboxes'] = result_instseg['bboxes']
        result['pred_scores'] = result_instseg['scores']
        # 3D
        result['pred_sRT'] = f_sRT
        result['pred_s'] = f_s
        result['pred_R'] = f_R
        result['pred_t'] = f_t
        result['pred_scales'] = f_size
        result['obj_pcs'] = obj_pcs

        img_posed = copy.deepcopy(img)
        img_posed = draw_pose_detections(img_posed, self.cam_K, result['pred_sRT'], result['pred_scales'], result_instseg['class_ids'])
        img_posed_msg = convert_nparray_to_Image(img_posed)

        return result, img_posed_msg

    def run_DPDN(self, image_raw:Image, depth_raw:Image, result_instseg):
        if result_instseg is None:
            return None, image_raw
                
        img = convert_Image_to_nparray(image_raw)
        depth = convert_Image_to_nparray(depth_raw)

        valid_inst, [f_points, f_rgb, f_choose, f_catId, f_prior] = self.prepare_cropped_data(img, depth, result_instseg)

        num_valid_inst = len(valid_inst)
        if num_valid_inst == 0:
            return None, image_raw
        valid_inst_np = np.array(valid_inst)
        result_instseg['class_ids'] = result_instseg['class_ids'][valid_inst_np]
        result_instseg['bboxes'] = result_instseg['bboxes'][valid_inst_np]
        result_instseg['scores'] = result_instseg['scores'][valid_inst_np]

        if len(valid_inst):
            pts = torch.FloatTensor(f_points).to(self.device)[valid_inst_np]
            rgb = torch.stack(f_rgb, dim=0).to(self.device)[valid_inst_np]
            choose = torch.LongTensor(f_choose).to(self.device)[valid_inst_np]
            label = torch.LongTensor(f_catId).to(self.device)[valid_inst_np]
            proir = torch.FloatTensor(f_prior).to(self.device)[valid_inst_np]

            inputs = {
                'rgb': rgb,
                'pts': pts,
                'choose': choose,
                'prior': proir,
                'category_label': label                
            }
            # inference
            end_points = self.objpose_model(inputs)

            pred_translation = end_points['pred_translation']
            pred_size = end_points['pred_size']
            pred_scale = torch.norm(pred_size, dim=1, keepdim=True)
            pred_size = pred_size / pred_scale
            pred_rotation = end_points['pred_rotation']
            pred_fullpc = end_points['pred_qv']

            num_instance = pred_rotation.size(0)
            pred_RTs =torch.eye(4).unsqueeze(0).repeat(num_instance, 1, 1).float().to(pred_rotation.device)
            pred_RTs[:, :3, 3] = pred_translation
            pred_RTs[:, :3, :3] = pred_rotation * pred_scale.unsqueeze(2)
            pred_scales = pred_size

            pred_i_sR = pred_rotation / pred_scale.unsqueeze(2)
            cuda_obj_pcs = torch.bmm(torch.add(torch.cuda.FloatTensor(f_points), -pred_translation.unsqueeze(1)), pred_i_sR)

            # filter point cloud in OBJCOORD
            obj_pcs = np.zeros_like(f_points)
            for i in range(cuda_obj_pcs.shape[0]):
                obj_pc = cuda_obj_pcs[i].detach().cpu().numpy()
                assert obj_pc.shape[1] == 3
                norm_pt = np.linalg.norm(obj_pc, axis=1)
                index = np.where(norm_pt < 0.5)
                obj_pc_filterd = obj_pc[index]
                if len(obj_pc_filterd) >= self.n_pts:
                    sample_index = np.random.choice(obj_pc_filterd.shape[0], self.n_pts, replace=False)
                else:
                    sample_index = np.random.choice(obj_pc_filterd.shape[0], self.n_pts, replace=True)
                    obj_pc_filterd = obj_pc_filterd[sample_index, :]
                obj_pcs[i] = obj_pc_filterd

                if self.show_pc:
                    import trimesh
                    pc = trimesh.PointCloud(obj_pc, colors=np.array([255,0,0,255]))
                    pc_filtered = trimesh.PointCloud(obj_pc_filterd)
                    pred_pc = trimesh.PointCloud(pred_fullpc[i].detach().cpu().numpy(), colors=np.array([0,255,0,255]))
                    trimesh.Scene([pc, pc_filtered, pred_pc]).show()                    
        result = {}
        # 2D
        result['pred_class_ids'] = result_instseg['class_ids']
        result['pred_bboxes'] = result_instseg['bboxes']
        result['pred_scores'] = result_instseg['scores']
        # 3D
        result['pred_sRT'] = pred_RTs.detach().cpu().numpy()
        result['pred_s'] = pred_scale.detach().cpu().numpy()
        result['pred_R'] = pred_rotation.detach().cpu().numpy()
        result['pred_t'] = pred_translation.detach().cpu().numpy()
        result['pred_scales'] = pred_scales.detach().cpu().numpy()
        result['obj_pcs'] = obj_pcs

        img_posed = copy.deepcopy(img)
        img_posed = draw_pose_detections(img_posed, self.cam_K, result['pred_sRT'], result['pred_scales'], result_instseg['class_ids'])
        img_posed_msg = convert_nparray_to_Image(img_posed)

        return result, img_posed_msg